import pandas as pd
import os

def process_and_collect_data(input_filename: str, url_param_name: str) -> pd.DataFrame | None:
    """
    处理单个分类CSV文件，生成下载链接，并返回包含所有原始列和新'下载链接'列的DataFrame。

    Args:
        input_filename (str): 输入CSV文件的名称。
        url_param_name (str): URL中ID参数的名称（例如 'ssid' 或 'sid'）。

    Returns:
        pd.DataFrame | None: 包含处理后数据的DataFrame，如果文件未找到或处理失败则返回None。
    """
    if not os.path.exists(input_filename):
        print(f"错误：文件 '{input_filename}' 未找到。请确保文件存在于脚本同一目录下。")
        return None

    try:
        df = pd.read_csv(input_filename, encoding='utf-8-sig')
    except Exception as e:
        print(f"读取文件 '{input_filename}' 时发生错误：{e}")
        return None

    def generate_download_url(row):
        base_url = "https://data.justeasy.cn/home/test/getModelCategoryList.php"
        return f"{base_url}?{url_param_name}={row['id']}&pagesize=3000"

    # 添加新列
    df['下载链接'] = df.apply(generate_download_url, axis=1)

    # --- 关键修改在这里 ---
    # 返回包含所有原始列和新添加的'下载链接'列的DataFrame
    return df

def main():
    """
    主函数，协调数据处理和合并流程。
    """
    print("--- 开始数据处理 ---")

    # 处理 ssid_categories.csv
    print("\n处理 ssid_categories.csv...")
    # df_ssid = process_and_collect_data('GZ_Store.csv', 'ssid')
    df_ssid = None
    # 处理 sid_categories.csv
    print("\n处理 sid_categories.csv...")
    df_sid = process_and_collect_data('GZsid_categories.csv', 'sid')
    # df_sid = None

    # --- 合并处理后的 DataFrames ---
    merged_df = pd.DataFrame()

    if df_ssid is not None:
        merged_df = pd.concat([merged_df, df_ssid], ignore_index=True)
        print("成功处理并添加了来自 'ssid_categories.csv' 的数据。")
    else:
        print("未从 'ssid_categories.csv' 收集到有效数据。")

    if df_sid is not None:
        merged_df = pd.concat([merged_df, df_sid], ignore_index=True)
        print("成功处理并添加了来自 'sid_categories.csv' 的数据。")
    else:
        print("未从 'sid_categories.csv' 收集到有效数据。")

    # --- 保存合并结果 ---
    output_file_merged = 'GZsid_with_links.csv'

    if not merged_df.empty:
        try:
            # 确保输出目录存在
            output_dir = os.path.dirname(output_file_merged)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir)
                print(f"创建输出目录: {output_dir}")

            merged_df.to_csv(output_file_merged, index=False, encoding='utf-8-sig')
            print(f"\n处理完成！所有收集到的数据已合并并保存到 {output_file_merged}")
            print("\n这是合并数据的预览（前5行）：")
            print(merged_df.head(5))
        except Exception as e:
            print(f"保存合并文件 '{output_file_merged}' 时发生错误：{e}")
    else:
        print("\n没有成功处理任何数据。没有创建合并输出文件。")

# 当脚本直接运行时，执行 main 函数
if __name__ == "__main__":
    main()